import os
import numpy as np
from PIL import Image

'''标签映射：19类映射为11类'''
data_root = 'E:\change_data\c_data_0905key2'

# cnn_to_ours_conversion = [1, 4, 2, 6, 9, 3, 7, 10, 11, 5, 8, 0]  # 旧的 1-11 到新的 1-11
# num_labels = 12
# todo  cityspace 模式下
cnn_to_ours_conversion = [3, 5, 1, 10, 10, 11, 9, 9, 4, 4, 2, 7, 8, 6, 6, 6, 6, 8, 8, 0]  # 旧的 0-19 到新的 0-11
num_labels = 19
num_labels_new = 11
num_img = len([f for f in os.listdir(os.path.join(data_root, 'label_2d_bin_raw')) if f.endswith('.bin')])

for frame_id in range(num_img):
    print(f'Processing {frame_id + 1} out of {num_img}')

    # 更改标签图像
    raw_cnn_img_file = os.path.join(data_root, 'label_2d_img_raw', f'{frame_id:06d}_bw.png')
    new_cnn_img_file = os.path.join(data_root, 'label_2d_img', f'{frame_id:06d}_bw.png')

    raw_cnn_img = Image.open(raw_cnn_img_file)
    raw_cnn_img = np.array(raw_cnn_img, dtype=np.uint8)
    # raw_cnn_img = raw_cnn_img + 1  # 0-10 到 1-11

    for row in range(raw_cnn_img.shape[0]):
        for col in range(raw_cnn_img.shape[1]):
            # raw_cnn_img[row, col] = cnn_to_ours_conversion[raw_cnn_img[row, col] - 1] - 1  # 1-11 到 0-10
            raw_cnn_img[row, col] = cnn_to_ours_conversion[raw_cnn_img[row, col]]  # 1-11 到 0-10

    raw_cnn_img = Image.fromarray(raw_cnn_img.astype(np.uint8))
    raw_cnn_img.save(new_cnn_img_file)

    img_height, img_width = raw_cnn_img.height, raw_cnn_img.width

    # 更改标签二进制文件
    raw_cnn_bin_file = os.path.join(data_root, 'label_2d_bin_raw', f'{frame_id:06d}.bin')
    new_cnn_bin_file = os.path.join(data_root, 'label_2d_bin', f'{frame_id:06d}.bin')

    with open(raw_cnn_bin_file, 'rb') as f:
        U = np.fromfile(f, dtype=np.float32)  # 480*640*19

    prob = U.reshape((img_height * img_width, num_labels))  # 307200*19

    # 重新映射概率
    # newprob = np.zeros_like(prob)
    newprob = np.zeros((img_height * img_width, num_labels_new))
    # for class_idx in range(num_labels):
    #     newprob[:, cnn_to_ours_conversion[class_idx]] = prob[:, class_idx]
    for class_idx in range(num_labels):  # todo 少一个other类后续检查是否有隐患
        new_class_idx = cnn_to_ours_conversion[class_idx]
        if newprob[:, new_class_idx - 1].any():  # 如果该新标签已经有值，证明多个老标签映射到了同一个新标签
            newprob[:, new_class_idx - 1] += prob[:, class_idx]
        else:
            newprob[:, new_class_idx - 1] = prob[:, class_idx]

    with open(new_cnn_bin_file, 'wb') as f:
        # newprob.T.astype(np.float32).tofile(f)
        newprob.astype(np.float32).tofile(f)

